defmodule Electric.ShapeCleanerTest do
  use ExUnit.Case, async: true
  use Repatch.ExUnit, assert_expectations: true

  import ExUnit.CaptureLog
  import Support.ComponentSetup
  import Support.TestUtils

  alias Electric.ShapeCache
  alias Electric.ShapeCache.Storage
  alias Electric.ShapeCache.ShapeCleaner
  alias Electric.Replication.LogOffset

  @stub_inspector Support.StubInspector.new(
                    tables: [{1, {"public", "items"}}],
                    columns: [
                      %{name: "id", type: "text", type_id: {25, 1}, is_generated: false},
                      %{name: "value", type: "text", type_id: {25, 1}, is_generated: false}
                    ]
                  )

  @shape Electric.Shapes.Shape.new!("items", inspector: @stub_inspector)
  @zero_offset LogOffset.last_before_real_offsets()
  @pg_snapshot_xmin_10 {10, 11, [10]}

  @moduletag :tmp_dir

  # Provide an inspector for downstream setup helpers (shape log collector, etc.)
  setup do
    %{inspector: @stub_inspector, pool: nil}
  end

  setup [
    :with_persistent_kv,
    :with_stack_id_from_test,
    :with_async_deleter,
    :with_pure_file_storage,
    :with_shape_status,
    :with_lsn_tracker,
    :with_status_monitor,
    :with_shape_cleaner
  ]

  for {cleanup_fn, suite_title} <- [
        {&ShapeCleaner.remove_shape/2, "remove_shape/2"},
        {&ShapeCleaner.remove_shape_async/2, "remove_shape_async/2"}
      ] do
    @cleanup_fn cleanup_fn

    describe "#{suite_title}" do
      setup(ctx) do
        %{stack_id: stack_id} = ctx
        existing_shapes = Map.get(ctx, :restore_shapes, [])

        patch_shape_status(list_shapes: fn ^stack_id -> existing_shapes end)

        Support.TestUtils.activate_mocks_for_descendant_procs(
          Electric.Replication.ShapeLogCollector
        )

        parent = self()

        patch_calls(
          Electric.StatusMonitor,
          mark_shape_log_collector_ready: fn _, _ ->
            send(parent, :shape_log_collector_ready)
            :ok
          end
        )

        existing_shape_map = Map.new(existing_shapes)

        if map_size(existing_shape_map) > 0 do
          patch_calls(Electric.ShapeCache.ShapeStatus,
            remove_shape: fn _stack_id, handle ->
              if shape = Map.get(existing_shape_map, handle) do
                {:ok, shape}
              else
                {:error, "invalid shape #{handle}"}
              end
            end
          )
        end

        :ok
      end

      setup [
        :with_log_chunking,
        :with_registry,
        :with_shape_log_collector,
        :with_noop_publication_manager
      ]

      setup(ctx) do
        Electric.Replication.ShapeLogCollector.mark_as_ready(ctx.stack_id)

        assert_receive :shape_log_collector_ready, 1000

        # don't suspend consumers in this test -- we don want them terminating
        # themselves before we clean them up
        Electric.StackConfig.put(ctx.stack_id, :shape_enable_suspend?, false)

        :ok
      end

      test "cleans up shape data and rotates the shape handle", ctx do
        Support.TestUtils.patch_snapshotter(fn parent,
                                               shape_handle,
                                               _shape,
                                               %{storage: storage} ->
          GenServer.cast(parent, {:pg_snapshot_known, shape_handle, @pg_snapshot_xmin_10})
          Storage.make_new_snapshot!([["test"]], storage)
          GenServer.cast(parent, {:snapshot_started, shape_handle})
        end)

        with_shape_cache(ctx)

        {shape_handle, _} = ShapeCache.get_or_create_shape_handle(@shape, ctx.stack_id)
        assert :started = ShapeCache.await_snapshot_start(shape_handle, ctx.stack_id)

        consumer_ref =
          Electric.Shapes.Consumer.whereis(ctx.stack_id, shape_handle)
          |> Process.monitor()

        {_, storage_opts} = ctx.storage
        storage = Storage.for_shape(shape_handle, ctx.storage)
        writer = Storage.init_writer!(storage, @shape)

        Storage.append_to_log!(
          changes_to_log_items([
            %Electric.Replication.Changes.NewRecord{
              relation: {"public", "items"},
              record: %{"id" => "1", "value" => "Alice"},
              log_offset: LogOffset.new(Electric.Postgres.Lsn.from_integer(1000), 0)
            }
          ]),
          writer
        )

        assert Storage.snapshot_started?(storage)
        assert File.exists?(Path.join(storage_opts.base_path, shape_handle))

        assert Enum.count(Storage.get_log_stream(LogOffset.last_before_real_offsets(), storage)) ==
                 1

        registry_ref = make_ref()
        Registry.register(ctx.registry, shape_handle, registry_ref)

        expect_calls(Electric.Replication.PublicationManager,
          remove_shape: fn _, ^shape_handle -> :ok end
        )

        expect_calls(Electric.Replication.ShapeLogCollector,
          remove_shape: fn _, ^shape_handle -> :ok end
        )

        :ok = @cleanup_fn.(ctx.stack_id, shape_handle)

        assert_shape_cleanup(shape_handle)

        assert_receive {:DOWN, ^consumer_ref, :process, _pid, {:shutdown, :cleanup}}

        assert_receive {^registry_ref, :shape_rotation}

        refute File.exists?(Path.join(storage_opts.base_path, shape_handle))

        {shape_handle2, _} = ShapeCache.get_or_create_shape_handle(@shape, ctx.stack_id)

        assert :started = ShapeCache.await_snapshot_start(shape_handle2, ctx.stack_id)
        assert shape_handle != shape_handle2
      end

      @tag restore_shapes: [{"my-shape", @shape}]
      test "removes shape handle from shape log collector if no consumer runnning", ctx do
        assert ["my-shape"] = Electric.Replication.ShapeLogCollector.active_shapes(ctx.stack_id)

        :ok = ShapeCleaner.remove_shape(ctx.stack_id, "my-shape")

        assert_shape_cleanup("my-shape")
        assert_shape_log_collector_active_shapes(ctx, [])
      end

      test "remove_shape swallows error if no shape to clean up", ctx do
        shape_handle = "foo"

        Support.TestUtils.patch_snapshotter(fn parent,
                                               shape_handle,
                                               _shape,
                                               %{storage: storage} ->
          GenServer.cast(parent, {:pg_snapshot_known, shape_handle, @pg_snapshot_xmin_10})
          Storage.make_new_snapshot!([["test"]], storage)
          GenServer.cast(parent, {:snapshot_started, shape_handle})
        end)

        with_shape_cache(ctx)

        {:ok, _} =
          with_log(fn -> ShapeCleaner.remove_shape(ctx.stack_id, shape_handle) end)
      end
    end
  end

  describe "remove_shapes_for_relations/2" do
    setup [
      :with_log_chunking,
      :with_registry,
      :with_shape_log_collector,
      :with_noop_publication_manager
    ]

    setup ctx do
      Support.TestUtils.patch_snapshotter(fn parent, shape_handle, _shape, %{storage: storage} ->
        GenServer.cast(parent, {:pg_snapshot_known, shape_handle, @pg_snapshot_xmin_10})
        Storage.make_new_snapshot!([["test"]], storage)
        GenServer.cast(parent, {:snapshot_started, shape_handle})
      end)

      with_shape_cache(ctx)
    end

    test "cleans up shape data for relevant shapes", ctx do
      {shape_handle, _} = ShapeCache.get_or_create_shape_handle(@shape, ctx.stack_id)
      assert :started = ShapeCache.await_snapshot_start(shape_handle, ctx.stack_id)

      consumer_ref =
        Electric.Shapes.Consumer.whereis(ctx.stack_id, shape_handle)
        |> Process.monitor()

      storage = Storage.for_shape(shape_handle, ctx.storage)
      writer = Storage.init_writer!(storage, @shape)

      Storage.append_to_log!(
        changes_to_log_items([
          %Electric.Replication.Changes.NewRecord{
            relation: {"public", "items"},
            record: %{"id" => "1", "value" => "Alice"},
            log_offset: LogOffset.new(Electric.Postgres.Lsn.from_integer(1000), 0)
          }
        ]),
        writer
      )

      assert Storage.snapshot_started?(storage)

      assert Enum.count(Storage.get_log_stream(LogOffset.last_before_real_offsets(), storage)) ==
               1

      # Cleaning unrelated relations should not affect the shape
      :ok =
        ShapeCleaner.remove_shapes_for_relations(
          ctx.stack_id,
          [{@shape.root_table_id + 1, {"public", "different"}}]
        )

      refute_receive {:DOWN, ^consumer_ref, :process, _pid, {:shutdown, :cleanup}}, 100

      # Shouldn't raise
      assert :ok = Stream.run(Storage.get_log_stream(@zero_offset, storage))

      :ok =
        ShapeCleaner.remove_shapes_for_relations(
          ctx.stack_id,
          [{@shape.root_table_id, {"public", "items"}}]
        )

      assert_shape_cleanup(shape_handle)

      # Allow asynchronous queued removal to complete
      assert_receive {:DOWN, ^consumer_ref, :process, _pid, {:shutdown, :cleanup}}, 1_000

      {shape_handle2, _} = ShapeCache.get_or_create_shape_handle(@shape, ctx.stack_id)
      assert :started = ShapeCache.await_snapshot_start(shape_handle2, ctx.stack_id)
      assert shape_handle != shape_handle2
    end
  end

  defp assert_shape_log_collector_active_shapes(ctx, shape_handles_active, timeout \\ 500) do
    assert shape_handles_active ==
             Electric.Replication.ShapeLogCollector.active_shapes(ctx.stack_id)
  rescue
    e in ExUnit.AssertionError ->
      if timeout <= 0 do
        reraise e, __STACKTRACE__
      else
        Process.sleep(50)
        assert_shape_log_collector_active_shapes(ctx, shape_handles_active, timeout - 50)
      end
  end
end
